#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/9/7 13:56
# @Author  : Duanxx
# @Site    : 
# @File    : locust2
# @Software: PyCharm

import json
import math
import random
from subprocess import Popen, PIPE

from locust import task, HttpUser, between, LoadTestShape
from logger import logger

class MyCustomShape(LoadTestShape):
    # 设置压测持续时长，单位秒
    time_limit = 60
    # 每秒启动/停止用户数
    spawn_rate = 10

    def tick(self):
        """
        返回一个元组，包含两值：
            user_count -- 总用户数
            spawn_rate -- 每秒启动/停止用户数
        返回None时，停止负载测试
        """
        # 获取压测执行的时间
        run_time = self.get_run_time()

        # 运行时长在压测最大时长内，则继续执行
        if run_time < self.time_limit:
            user_count = round(run_time, -1)
            return user_count, self.spawn_rate
        return None


class StepShape(LoadTestShape):
    """
        step_time -- 步骤之间的时间
        step_load -- 用户在每一步增加数量
        spawn_rate -- 用户在每一步式每秒停止/启动数量
        time_limit -- 时间限制，以秒为单位
    """

    step_time = 30
    step_load = 10
    spawn_rate = 10
    time_limit = 600

    def tick(self):
        run_time = self.get_run_time()

        if run_time > self.time_limit:
            return None

        current_step = math.floor(run_time / self.step_time) + 1
        return current_step * self.step_load, self.spawn_rate


class TimeShape(LoadTestShape):
    """
        duration -- 多少秒后进入下一个阶段
        users -- 用户数
        spawn_rate -- 每秒要启动/停止的用户数
    """

    stages = [
        {"duration": 600, "users": 1, "spawn_rate": 1},
        {"duration": 600, "users": 5, "spawn_rate": 1},
        {"duration": 600, "users": 8, "spawn_rate": 1},
        {"duration": 600, "users": 10, "spawn_rate": 1},
        {"duration": 600, "users": 12, "spawn_rate": 1},
        {"duration": 600, "users": 15, "spawn_rate": 1},
        {"duration": 600, "users": 20, "spawn_rate": 1},
        {"duration": 600, "users": 30, "spawn_rate": 1},
        {"duration": 600, "users": 40, "spawn_rate": 1},
        {"duration": 600, "users": 50, "spawn_rate": 1},
        {"duration": 600, "users": 70, "spawn_rate": 1},
        {"duration": 600, "users": 100, "spawn_rate": 1}
    ]

    def tick(self):
        run_time = self.get_run_time()

        for stage in self.stages:
            if run_time < stage["duration"]:
                tick_data = (stage["users"], stage["spawn_rate"])
                return tick_data

        return None
